import cv2 as cv
import numpy as np

def interpolate_frame(frame1, frame2, alpha):
    """高质量帧插值函数，解决黑边问题"""
    # 转为灰度计算光流
    gray1 = cv.cvtColor(frame1, cv.COLOR_BGR2GRAY)
    gray2 = cv.cvtColor(frame2, cv.COLOR_BGR2GRAY)

    # 使用更精确的光流算法
    dis = cv.DISOpticalFlow_create(cv.DISOPTICAL_FLOW_PRESET_MEDIUM)
    flow_fw = dis.calc(gray1, gray2, None)
    flow_bw = dis.calc(gray2, gray1, None)

    h, w = frame1.shape[:2]
    y, x = np.mgrid[0:h, 0:w].astype(np.float32)

    # 正向映射（添加边界保护）
    map_x1 = np.clip(x + alpha * flow_fw[..., 0], 0, w - 1)
    map_y1 = np.clip(y + alpha * flow_fw[..., 1], 0, h - 1)

    # 反向映射（添加边界保护）
    map_x2 = np.clip(x - (1 - alpha) * flow_bw[..., 0], 0, w - 1)
    map_y2 = np.clip(y - (1 - alpha) * flow_bw[..., 1], 0, h - 1)

    # 使用镜像边界模式消除黑边
    warp1 = cv.remap(frame1, map_x1, map_y1, cv.INTER_LINEAR, borderMode=cv.BORDER_REFLECT)
    warp2 = cv.remap(frame2, map_x2, map_y2, cv.INTER_LINEAR, borderMode=cv.BORDER_REFLECT)

    # 安全混合（避免溢出）
    return cv.addWeighted(warp1, 1 - alpha, warp2, alpha, 0)


if __name__ == '__main__':
    input_path = input("输入补帧目标路径：")
    output_path = input("存储路径:") # 修改输出文件名

    cap = cv.VideoCapture(input_path)
    input_fps = cap.get(cv.CAP_PROP_FPS)  # 获取原始帧率（假设30fps）
    target_fps =input_fps*2   # 目标帧率

    # 检查视频是否打开成功
    if not cap.isOpened():
        print("无法打开输入视频！")
        exit()

    ret, prev_frame = cap.read()
    if not ret:
        print("无法读取第一帧！")
        exit()

    # 视频写入器设置（60fps）
    fourcc = cv.VideoWriter_fourcc(*'mp4v')
    video_writer = cv.VideoWriter(
        output_path,
        fourcc,
        target_fps,  # 关键修改：设置为60fps
        (int(cap.get(cv.CAP_PROP_FRAME_WIDTH)),
         int(cap.get(cv.CAP_PROP_FRAME_HEIGHT)))
    )

    frame_count = 0
    while True:
        ret, next_frame = cap.read()
        if not ret:
            break

        # 写入原始帧（保持时间点准确）
        video_writer.write(prev_frame)
        frame_count += 1

        # 生成并写入中间帧（实现60fps）
        mid_frame = interpolate_frame(prev_frame, next_frame, alpha=0.5)
        video_writer.write(mid_frame)
        frame_count += 1

        # 显示进度
        print(f"处理进度: {frame_count // 2}/{int(cap.get(cv.CAP_PROP_FRAME_COUNT))} 帧", end='\r')
        prev_frame = next_frame

    # 释放资源
    cap.release()
    video_writer.release()
    cv.destroyAllWindows()

    print(f"\n补帧完成！原始帧率：{input_fps}fps → 输出帧率：{target_fps}fps")
    print(f"视频已保存至：{output_path}")


